# -*- coding: utf-8 -*-
# Created Time : 2022/1/5 15:19
# Author:Zhou
import os

import h5py
import matplotlib.pyplot as plt
import numpy as np
from ANC_tools import sin_wave, awgn, NPP_1, NPP_2, NPP_3, NPP_4, NPP_5, logisticChaotic
import scipy.io as scio
import win32com.client

# some settings
# 生成训练数据
# mix是输入
# 对于speech enhancement，mix= noise + speech
# 对于ANC，mix = reference signal
# sph是目标
# 对于SE,sph是纯净的人声
# 对于ANC，sph是经过非线性过程的初级噪声

dataFileList = ['C:/Users/francklinson/PycharmProjects/ANC_PY/NoiseData/volvo_norm.mat',
                'C:/Users/francklinson/PycharmProjects/ANC_PY/NoiseData/factory1_norm.mat',
                'C:/Users/francklinson/PycharmProjects/ANC_PY/NoiseData/destroyerengine_norm.mat',
                'C:/Users/francklinson/PycharmProjects/ANC_PY/NoiseData/babble_norm.mat',
                'C:/Users/francklinson/PycharmProjects/ANC_PY/NoiseData/buccaneer1_norm.mat',
                'C:/Users/francklinson/PycharmProjects/ANC_PY/NoiseData/f16_norm.mat',
                'C:/Users/francklinson/PycharmProjects/ANC_PY/NoiseData/machinegun_norm.mat',
                'C:/Users/francklinson/PycharmProjects/ANC_PY/NoiseData/leopard_norm.mat']

NPP = NPP_4
rms = 1.0
n_tr_ex = 128
samplingRate = 16000
noiseLength = 3  # 时间 单位s
# mix = logisticChaotic(0.9, 4, 16000)

idx = 0
for dataFile in dataFileList:
    pass
    data = scio.loadmat(dataFile)
    noiseSource = data['data']
    for i in range(n_tr_ex):  # n_tr_ex is the number of training examples
        # generate a noisy mixture
        # mix stores a noisy mixture utterance
        # sph the corresponding clean speech utterance.

        # speech enhancement 的任务是输入mix(speech+noise)，得到接近sph的声音
        # ANC的任务是  输入reference (x(n)，单频参考噪声), 得到接近初级噪声(P(n))的信号
        print('generating: sample_{}'.format(idx))
        beginning = np.random.randint(0, len(noiseSource) - noiseLength * samplingRate)
        end = beginning + noiseLength * samplingRate
        mix = noiseSource[beginning: end, :]
        sph = NPP(mix)
        # normalize
        c = rms * np.sqrt(mix.size / np.sum(mix ** 2))
        mix *= c
        sph *= c
        # plt.figure()
        # plt.subplot(2, 1, 1)
        # plt.plot(mix)
        # plt.subplot(2, 1, 2)
        # plt.plot(sph)
        # plt.show()

        # save file
        filepath = 'datasets/hybrid_NPP4/tr/'
        filename = 'tr_{}.ex'.format(idx)
        writer = h5py.File(os.path.join(filepath, filename), 'w')
        writer.create_dataset('mix', data=mix.astype(np.float32), shape=mix.shape, chunks=True)
        writer.create_dataset('sph', data=sph.astype(np.float32), shape=sph.shape, chunks=True)
        writer.close()
        idx += 1

for i in range(n_tr_ex*8):
    print("../data/datasets/hybrid_NPP4/tr/tr_{}.ex".format(i))
speak = win32com.client.Dispatch("SAPI.SpVoice")
# 创建发声对象
speak.Speak('数据集构造完成！')
